/*
 *
 *  * Copyright 2020 New Relic Corporation. All rights reserved.
 *  * SPDX-License-Identifier: Apache-2.0
 *
 */

package com.newrelic.agent.utilization;

import com.newrelic.agent.Agent;
import com.newrelic.agent.MetricNames;
import org.json.simple.JSONObject;
import org.json.simple.parser.JSONParser;

import java.text.MessageFormat;
import java.util.HashMap;
import java.util.Map;
import java.util.logging.Level;

/**
 * Use the AWS Instance Metadata Service (IMDS) v2 to get instance data for utilization purposes. See
 * <a href="https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html">AWS IMDS documentation</a>
 * for more information.
 */
public class AWS implements CloudVendor {
    static String PROVIDER = "aws";
    private final CloudUtility cloudUtility;

    public AWS(CloudUtility cloudUtility) {
        this.cloudUtility = cloudUtility;
    }

    private static final String INSTANCE_TOKEN_URL = "http://169.254.169.254/latest/api/token";
    private static final String INSTANCE_DOCUMENT_URL = "http://169.254.169.254/2016-09-02/dynamic/instance-identity/document";

    private static final int REQUEST_TIMEOUT_MILLIS = 100;
    private static final int TOKEN_TTL_SECONDS = 60;

    // AWS request strings
    private static final String AWS_INSTANCE_ID_REQUEST = "instanceId";
    private static final String AWS_INSTANCE_TYPE_REQUEST = "instanceType";
    private static final String AWS_AVAILABILITY_ZONE_REQUEST = "availabilityZone";

    // AWS map keys. These are the keys that will be added to the vendor hash in the JSON generated by the agent.
    private static final String AWS_INSTANCE_ID_KEY = "instanceId";
    private static final String AWS_INSTANCE_TYPE_KEY = "instanceType";
    private static final String AWS_AVAILABILITY_ZONE_KEY = "availabilityZone";

    /**
     * Query the AWS API to get metadata
     *
     * @return {@link AwsData} with values for instanceId, instanceType, and availabilityZone
     */
    @Override
    public AwsData getData() {
        String token = getAwsToken();

        if (token == null || token.isEmpty()) {
            return AwsData.EMPTY_DATA;
        }

        try {
            String unparsedResult = getAwsValues(token);

            JSONParser parser = new JSONParser();
            JSONObject result = null;

            if (unparsedResult != null) {
                result = (JSONObject) parser.parse(unparsedResult);
            }

            // not on AWS
            if (result == null || result.isEmpty()) {
                return AwsData.EMPTY_DATA;
            }

            String type = (String) result.get(AWS_INSTANCE_TYPE_REQUEST);
            String id = (String) result.get(AWS_INSTANCE_ID_REQUEST);
            String zone = (String) result.get(AWS_AVAILABILITY_ZONE_REQUEST);

            if (cloudUtility.isInvalidValue(type) || cloudUtility.isInvalidValue(id)
                    || cloudUtility.isInvalidValue(zone)) {
                Agent.LOG.log(Level.WARNING, "Failed to validate AWS value");
                recordAwsError();
                return AwsData.EMPTY_DATA;
            }

            AwsData data = new AwsData(id, type, zone);
            Agent.LOG.log(Level.FINEST, "Found {0}", data);
            return data;
        } catch (Exception e) {
            return AwsData.EMPTY_DATA;
        }
    }

    private String getAwsToken() {
        try {
            return cloudUtility.httpPut(INSTANCE_TOKEN_URL, REQUEST_TIMEOUT_MILLIS, "X-aws-ec2-metadata-token-ttl-seconds: " + TOKEN_TTL_SECONDS);
        } catch (Throwable t) {
            Agent.LOG.log(Level.FINEST, t, "Error occurred trying to get AWS token");
            recordAwsError();
        }
        return null;
    }

    private String getAwsValues(String token) {
        try {
            return cloudUtility.httpGet(INSTANCE_DOCUMENT_URL, REQUEST_TIMEOUT_MILLIS, "X-aws-ec2-metadata-token: " + token);
        } catch (Throwable t) {
            Agent.LOG.log(Level.FINEST, MessageFormat.format("Error occurred trying to get AWS values. {0}", t));
            recordAwsError();
        }
        return null;
    }

    private void recordAwsError() {
        cloudUtility.recordError(MetricNames.SUPPORTABILITY_AWS_ERROR);
    }

    protected static class AwsData implements CloudData {
        private final String instanceId;
        private final String instanceType;
        private final String availabilityZone;

        static final AwsData EMPTY_DATA = new AwsData();

        private AwsData() {
            instanceId = null;
            instanceType = null;
            availabilityZone = null;
        }

        protected AwsData(String id, String type, String zone) {
            instanceId = id;
            instanceType = type;
            availabilityZone = zone;
        }

        public String getInstanceId() {
            return instanceId;
        }

        public String getInstanceType() {
            return instanceType;
        }

        public String getAvailabilityZone() {
            return availabilityZone;
        }

        @Override
        public Map<String, String> getValueMap() {
            Map<String, String> aws = new HashMap<>();

            if (instanceType == null || instanceId == null || availabilityZone == null) {
                return aws;
            } else {
                aws.put(AWS_INSTANCE_TYPE_KEY, instanceType);
                aws.put(AWS_INSTANCE_ID_KEY, instanceId);
                aws.put(AWS_AVAILABILITY_ZONE_KEY, availabilityZone);
            }
            return aws;
        }

        @Override
        public String getProvider() {
            return PROVIDER;
        }

        @Override
        public boolean isEmpty() {
            return this == EMPTY_DATA;
        }

        @Override
        public String toString() {
            return "AwsData{" +
                    "instanceId='" + instanceId + '\'' +
                    ", instanceType='" + instanceType + '\'' +
                    ", availabilityZone='" + availabilityZone + '\'' +
                    '}';
        }
    }

}
